package cn.denvie.agent.collector;

import cn.denvie.agent.utils.StringUtils;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.LoaderClassPath;
import javassist.NotFoundException;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Level;
import java.util.logging.Logger;

public abstract class AbstractByteTransformCollector {

    private static Logger logger = Logger.getLogger(AbstractByteTransformCollector.class.getName());
    private static Map<ClassLoader, ClassPool> classPoolMap = new ConcurrentHashMap<>();

    public AbstractByteTransformCollector(Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                                    ProtectionDomain protectionDomain, byte[] classfileBuffer) {
                if (loader == null)
                    return null;
                if (className == null)
                    return null;
                className = className.trim().replace("/", ".");
                try {
                    return AbstractByteTransformCollector.this.transform(loader, className);
                } catch (Exception e) {
                    logger.log(Level.SEVERE, "ClassFileTransformer插桩失败: ", e);
                }
                return null;
            }
        });
    }

    // 字节码插桩，由子类实现
    public abstract byte[] transform(ClassLoader loader, String className) throws Exception;

    protected static CtClass toCtClass(ClassLoader loader, String className) throws NotFoundException {
        if (!classPoolMap.containsKey(loader)) {
            ClassPool classPool = new ClassPool();
            classPool.insertClassPath(new LoaderClassPath(loader));
            classPoolMap.put(loader, classPool);
        }
        ClassPool cp = classPoolMap.get(loader);
        className = className.replaceAll("/", ".");
        return cp.get(className);
    }

    /**
     * 用逗号分隔为列表。
     */
    protected List<String> separateByComma(String text) {
        if (StringUtils.isEmpty(text)) {
            return Collections.emptyList();
        }
        String[] splits = text.split(",");
        return Arrays.asList(splits);
    }

}